# coding=utf-8

import sys
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import shuffle, randint, uniform

from icecream import ic

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)


class Node:
    def __init__(self, id_: int, x: float, y: float):
        self.id = id_
        self.x = x
        self.y = y
        self.left = None
        self.right = None


class Data:
    """ Read a TSPLIB file, and store it as a basic data. """

    def __init__(self, filepath: str):
        self.data = np.loadtxt(filepath, skiprows=6, encoding='utf-8', dtype=float)
        self.size = self.data.shape[0]  # size of city set
        self.distance = self.calc_distance()  # distance matrix
        self.nnList = self.calc_nnList()  # nearest neighbor list

    def calc_distance(self):
        i, x, y = self.data[:, 0].astype(int), self.data[:, 1], self.data[:, 2]
        i = np.insert(i, obj=0, values=0, axis=0)
        j = i.reshape(len(i), -1)
        x, y = np.insert(x, obj=0, values=0, axis=0), np.insert(y, obj=0, values=0, axis=0)
        distance = np.sqrt((x[i] - x[j]) ** 2 + (y[i] - y[j]) ** 2)
        distance[0, :], distance[:, 0] = 0, 0
        return distance

    def calc_nnList(self):
        distance = np.delete(self.distance, obj=0, axis=0)
        distance = np.delete(distance, obj=0, axis=1)
        nnList = np.argsort(distance, axis=1) + 1
        nnList = np.insert(nnList, obj=0, values=0, axis=0)
        return nnList

    def nearest_neighbor(self, node):
        """ return the id and distance of nearest neighbor of 'node' """
        return self.nnList[node][1], self.distance[node][self.nnList[node][1]]


class Individual:
    def __init__(self, _data: Data):
        self.tour = [Node(id_=data.data[i][0].astype(int), x=data.data[i][1], y=data.data[i][2])
                     for i in range(data.size)]  # basic tour consists of city nodes
        self.size = _data.size  # the size of a tour
        shuffle(self.tour)
        # self.two_opt()
        for i in range(self.size):  # mark left and right neighbors
            self.tour[i].left = self.tour[(i + self.size - 1) % self.size].id
            self.tour[i].right = self.tour[(i + 1) % self.size].id
        self.tour.sort(key=lambda _: _.id)  # sort node by id
        self.length = self.calc_length(_data.distance)

    def calc_length(self, distance):
        length = 0
        for i in range(self.size):
            length += distance[self.tour[i].id][self.tour[i].left]
            length += distance[self.tour[i].id][self.tour[i].right]
        length /= 2
        return length

    def two_opt(self, distance):
        improve = True
        while improve:
            improve = False
            break

    def show(self):
        tour = [[node.left, node.id, node.right] for node in self.tour]


def select_edge(trace, stock, curr):
    """ search next node of current node in stock
        @param trace: traced edges
        @param stock: stock edges of A or B
        @param curr: current node
        @return trace, stock, next node
    """
    # search node
    if uniform(0, 1) < 0.5:  # left first
        if stock[curr - 1].left != 0:
            _next = stock[curr - 1].left
        else:
            _next = stock[curr - 1].right
    else:  # right first
        if stock[curr - 1].right != 0:
            _next = stock[curr - 1].right
        else:
            _next = stock[curr - 1].left
    # add node
    trace.append(deepcopy(stock[_next - 1]))
    # remove the searched edge (bidirectionally)
    # cut link of 'curr -> next'
    if stock[curr - 1].left == _next:
        stock[curr - 1].left = 0
    else:
        stock[curr - 1].right = 0
    # cut link of 'next -> curr'
    if stock[_next - 1].left == curr:
        stock[_next - 1].left = 0
    else:
        stock[_next - 1].right = 0
    # if the left and right neighbors are both 0, set the node to 0
    if stock[curr - 1].left == stock[curr - 1].right:
        stock[curr - 1].id = 0
    if stock[_next - 1].left == stock[_next - 1].right:
        stock[_next - 1].id = 0
    return trace, stock, _next


def judge_AB_cycle(trace, ESet):
    end = len(trace) - 1  # the last of trace
    print(f'trace size: {end}')
    for i in range(end):
        if trace[i].id == trace[end].id & (end - i) % 2 == 0:  # AB-cycle detected
            print('AB-cycle detected')
            AB_cycle = deepcopy(trace[i+1:])
            trace = trace[:i+1]
            if end - i > 2:  # more than 2 edge
                ESet.append(AB_cycle)
            break
    return trace, ESet


def EAX(pA, pB):
    if pA == pB:
        return
    ESet = []
    stockA = deepcopy(pA.gene)  # edges not traced in A
    stockB = deepcopy(pB.gene)  # edges not traced in B
    size = pA.size
    trace = []
    while True:
        # random choose the first node
        while True:
            curr = randint(1, size)
            if stockA[curr - 1].id != 0:
                break
        group = 0
        # trace edges of tourA and tourB in turn
        trace.clear()
        trace.append(deepcopy(stockA[curr - 1]))
        print_stock(stockA, stockB)
        print_trace(trace)
        while True:
            if group == 0:
                trace, stockA, curr = select_edge(trace, stockA, curr)  # trace A
                group = 1
            else:
                trace, stockB, curr = select_edge(trace, stockB, curr)  # trace B
                group = 0

            print_stock(stockA, stockB)
            print_trace(trace)
            trace, ESet = judge_AB_cycle(trace, ESet)
            plot_map(plt, data, pA.gene, pB.gene)
            plot_trace(plt, trace)
            plt.show()
            print('-----------------------------')
            if len(trace) == 1:  # current trace is empty
                break
        pass


def print_stock(stockA, stockB):
    for i in range(len(stockA)):
        print(f'{stockA[i].left:2d} {stockA[i].id:2d} {stockA[i].right:2d} | '
              f'{stockB[i].left:2d} {stockB[i].id:2d} {stockB[i].right:2d}')


def print_trace(trace):
    for i in range(len(trace)):
        print(trace[i].left, trace[i].id, trace[i].right)


def plot_map(plot, _data, tourA, tourB):
    plot.scatter(_data.data[:, 1], _data.data[:, 2], c='k', s=20)
    for i in range(len(_data.data)):
        plot.text(x=_data.data[i, 1], y=_data.data[i, 2], s=_data.data[i, 0].astype(int))
    for node in tourA:
        plot.plot([tourA[node.left-1].x, node.x], [tourA[node.left-1].y, node.y], c='r', lw=1)
    # for node in tourB:
    #     plot.plot([tourB[node.left-1].x, node.x], [tourB[node.left-1].y, node.y], c='b', lw=1)
    return plot


def plot_trace(plot, trace):
    color = ['r', 'b']
    ls = ['-', '--']
    group = 1
    for i in range(1, len(trace)):
        group += 1
        plot.plot([trace[i-1].x,  trace[i].x], [trace[i-1].y,  trace[i].y], c=color[group % 2], ls=ls[group % 2], lw=2)
    return plot




TRIAL = 5
GENERATION = 100
NP = 10
NO = 5
if __name__ == '__main__':
    filename = f'../../data/TSPLIB/{sys.argv[1]}.tsp.txt'
    data = Data(filename)
    p = Individual(data)
    for trial in range(TRIAL):
        # init population
        population, offspring = [], []
        for n in range(NP):
            population.append(Individual(data))
        for generation in range(GENERATION):
            # select parents
            offspring.clear()
            for n in range(NP):
                parentA = population[n]
                parentB = population[(n + 1) % NP]
                # crossover
                offspring.append(EAX(parentA, parentB))
